{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# PocketFlow基础使用指南"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 1. 环境准备"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 安装依赖\n",
    "!pip install torch torchvision\n",
    "!pip install matplotlib"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 2. 加载基础模型"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import torchvision.models as models\n",
    "\n",
    "# 加载预训练ResNet18\n",
    "model = models.resnet18(pretrained=True)\n",
    "print(\"原始模型大小:\", sum(p.numel() for p in model.parameters()))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 3. 使用蒸馏压缩器"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from pocketflow.compressors import DistillationCompressor\n",
    "\n",
    "# 初始化教师模型(更大的模型)\n",
    "teacher = models.resnet50(pretrained=True)\n",
    "\n",
    "# 配置蒸馏压缩器\n",
    "config = {\n",
    "    'teacher_model': teacher,\n",
    "    'temperature': 4.0,\n",
    "    'alpha': 0.7,\n",
    "    'optimizer': {'type': 'Adam', 'lr': 1e-3}\n",
    "}\n",
    "\n",
    "# 创建压缩器实例\n",
    "compressor = DistillationCompressor(model, config)\n",
    "\n",
    "# 模拟训练循环\n",
    "for epoch in range(5):\n",
    "    # 这里使用伪数据\n",
    "    data = torch.randn(32, 3, 224, 224)\n",
    "    target = torch.randint(0, 1000, (32,))\n",
    "    \n",
    "    loss = compressor.train_step(data, target)\n",
    "    print(f\"Epoch {epoch}, Loss: {loss:.4f}\")\n",
    "\n",
    "# 保存蒸馏后模型\n",
    "compressor.save_model('distilled_model.pth')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 4. 使用混合精度压缩器"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from pocketflow.compressors import MixedPrecisionCompressor\n",
    "\n",
    "# 配置混合精度压缩器\n",
    "config = {\n",
    "    'enabled': True,\n",
    "    'opt_level': 'O2',\n",
    "    'loss_scale': 1024.0\n",
    "}\n",
    "\n",
    "# 创建压缩器实例\n",
    "compressor = MixedPrecisionCompressor(model, config)\n",
    "\n",
    "# 模拟训练步骤\n",
    "data = torch.randn(32, 3, 224, 224)\n",
    "target = torch.randint(0, 1000, (32,))\n",
    "loss = compressor.train_step(data, target)\n",
    "\n",
    "# 获取性能报告\n",
    "report = compressor.analyze_performance()\n",
    "print(\"性能报告:\", report)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 5. 使用量化压缩器"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from pocketflow.compressors import QuantizationCompressor\n",
    "\n",
    "# 配置量化压缩器\n",
    "config = {\n",
    "    'quant_type': 'static',\n",
    "    'bit_width': 8,\n",
    "    'qconfig': {\n",
    "        'activation': {'dtype': 'quint8'},\n",
    "        'weight': {'dtype': 'qint8'}\n",
    "    }\n",
    "}\n",
    "\n",
    "# 创建压缩器实例\n",
    "compressor = QuantizationCompressor(model, config)\n",
    "\n",
    "# 模拟校准过程\n",
    "for _ in range(5):\n",
    "    data = torch.randn(32, 3, 224, 224)\n",
    "    compressor.calibrate([(data, None)])\n",
    "\n",
    "# 应用量化\n",
    "compressor.quantize()\n",
    "\n",
    "# 保存量化模型\n",
    "compressor.save_model('quantized_model.pth')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 6. 结果可视化"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "\n",
    "# 模拟性能数据\n",
    "labels = ['原始模型', '蒸馏', '混合精度', '量化']\n",
    "size = [44.6, 44.6, 44.6, 11.2]  # MB\n",
    "speed = [100, 95, 180, 210]  # 推理速度(越高越好)\n",
    "\n",
    "fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(12, 4))\n",
    "\n",
    "ax1.bar(labels, size)\n",
    "ax1.set_title('模型大小比较(MB)')\n",
    "\n",
    "ax2.bar(labels, speed)\n",
    "ax2.set_title('推理速度比较(越高越好)')\n",
    "\n",
    "plt.tight_layout()\n",
    "plt.show()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}